import json
import os
import os.path as osp
import json
import cv2
import time
from cv2 import data
# import math
import numpy as np
import scipy.interpolate as spi

class MyEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(MyEncoder, self).default(obj)

class LabelDoor:
    def __init__(self, video_path=None):
        self.now_position = None  # 当前鼠标位置
        self.point_list = [None, None, None, None, None]  # 一条线的每个点的位置
        self.color = [(255, 0, 0), # 线
                (0, 0, 255),    # 线
                (0, 255, 0),    # 线
                (255, 255, 0),  # 线
                (0, 0, 0)]  # 笔刷

        self.frame_mode = 'Normal'  # Normal, Edit, AutoPlay
        self.edit_mode = 'none' # 编辑模式
        self.curve_mode = False  # 是否生成线
        self.bind_mode = False
        self.point_idx = -1  # 点的索引
        self.line_idx = -1  # 线的索引
        self.drawing = False  # 是否正在擦除
        self.m_erase = None  # 记录擦除的位置
        self.erase_size = 10
        self.last_im_success = None  # 上次成功加载的图像
        self.cap = None  # 视频流

        self.video_name = None  # 视频名
        self.video_path = video_path  # 视频路径
        if video_path:
            self.load_data(video_path)
        

    def label_mask(self, event, x, y, flags, param):
        self.now_position = (x, y)
        if event == cv2.EVENT_LBUTTONDOWN:  # 鼠标按下
            if self.point_idx >= 0 and self.point_idx <= 4:
                old_point = self.data['Line'][self.line_idx][self.point_idx]
                self.point_list[self.point_idx] = (x, y)
                if self.bind_mode and old_point is not None:  # 点绑定，同步移动
                    dx = x - old_point[0]
                    dy = y - old_point[1]
                    for i in range(5):
                        if i == self.point_idx:
                            continue
                        if self.data['Line'][self.line_idx][i] is not None:
                            nx = self.data['Line'][self.line_idx][i][0] + dx
                            ny = self.data['Line'][self.line_idx][i][1] + dy
                            self.data['Line'][self.line_idx][i] = (nx, ny)
            self.drawing = True
        elif event == cv2.EVENT_MOUSEMOVE:  # 鼠标移动
            pass
        elif event == cv2.EVENT_LBUTTONUP:  # 鼠标松开
            if self.point_idx >= 0 and self.point_idx <= 4:
                self.point_list[self.point_idx] = None
            self.drawing = False

    def draw_point(self, im, point_data, the_color, size=2):
        """ 画点 """
        x, y = point_data
        cv2.circle(im, (x, y), size, the_color, 2)

    def load_data(self, video_path):
        """ 加载数据 """
        assert osp.isfile(video_path)  # video/v.mp4
        (v_path, v_name) = os.path.split(video_path) # (video, v.mp4)
        (v_pre_name, _) = os.path.splitext(v_name)  # (v.mp4, _)
        # print(op.(video_path))
        self.video_name = v_name  # v.mp4
        base_path = osp.join(v_path, "..")
        json_dir = os.path.join(base_path, 'label')  # label/
        if not os.path.isdir(json_dir):
            os.makedirs(json_dir, exist_ok=True)
            
        
        self.data = {}
        self.data['Line'] = [[None, None, None, None, None],
                        [None, None, None, None, None],
                        [None, None, None, None, None],
                        [None, None, None, None, None]]
        self.data['h_samples'] = None
        self.data['lanes'] = []
        self.data['raw_file'] = None

        # file_name_json = v_pre_name + "_0001.json"
        # file_name_json = os.path.join(json_dir, file_name_json)  # label/v_1.json
        # if osp.isfile(file_name_json):
        #     # 打开文件并获取内容
        #     with open(file_name_json, 'r') as fp:
        #         self.data = json.load(fp)  # 加载json文件
        #         assert 'Line' in self.data and 'h_samples' in self.data and 'lanes' in self.data and 'raw_file' in self.data
        # else:
        #     self.data = {}
        #     self.data['Line'] = [[None, None, None, None, None],
        #                     [None, None, None, None, None],
        #                     [None, None, None, None, None],
        #                     [None, None, None, None, None]]
        #     self.data['h_samples'] = None
        #     self.data['lanes'] = []
        #     self.data['raw_file'] = None
        #     # 新建数据文件
        #     json.dump(self.data, open(file_name_json, 'w'), indent=4, cls=json.JSONEncoder)

    def save_data(self, pos_frame):
        self.data['h_samples'] = None
        self.data['lanes'] = []
        for idx in range(4):
            curve_x, curve_y = self.get_curve(self.data['Line'][idx], step=10)
            if curve_y is not None:
                self.data['h_samples'] = curve_y
                for i in range(len(curve_x)):
                    if self.m_erase[int(curve_x[i])][int(curve_y[i])]:
                        curve_x[i] = -1
                self.data['lanes'].append(curve_x)

        (v_path, v_name) = os.path.split(self.video_path) # (video, v.mp4)
        (v_pre_name, _) = os.path.splitext(v_name)  # (v, mp4)
        # print(op.(video_path))
        base_path = osp.join(v_path, "..", "label")  # label

        img_relative_path = osp.join('clips', '0520', v_pre_name) # clips/%Y%m/v  TODO
        img_path = osp.join(base_path, img_relative_path) # label/clips/%Y%m/v
        if not osp.isdir(img_path):
            os.makedirs(img_path, exist_ok=True)
        json_path = osp.join(base_path, 'single_label')  # label/single_label
        json_path = osp.join(json_path, v_pre_name)  # label/single_label/v
        if not osp.isdir(json_path):
            os.makedirs(json_path, exist_ok=True)
        file_name = v_pre_name + '_%04d'%pos_frame
        img_save_path = osp.join(img_path, file_name+'.jpg')  # label/clips/v_%d.jpg
        self.data['raw_file'] = osp.join(img_relative_path, file_name+'.jpg')
        json.dump(self.data, open(osp.join(json_path, file_name+'.json'), 'w'), indent=4, cls=MyEncoder)
        cv2.imwrite(img_save_path, self.last_im_success)
        print('保存成功！文件名：{}'.format(file_name))
    
    def merge_json(self, _del=True):
        # video_path: xxx/dir/v1.map4
        (v_path, v_name) = os.path.split(self.video_path) # (video, v.mp4)
        (v_pre_name, _) = os.path.splitext(v_name)  # (v, mp4)
        base_path = osp.join(v_path, "..", "label")  # label
        all_json_path = osp.join(base_path, v_pre_name+'.json')  # label/v.json
        jsons_path = osp.join(base_path, 'single_label')  
        jsons_path = osp.join(jsons_path, v_pre_name)  # label/single_label/v
        assert osp.isdir(jsons_path)

        json_list = []
        for sub_json_path in os.listdir(jsons_path):
            file_name_json = osp.join(jsons_path, sub_json_path)
            if osp.isfile(file_name_json):
                # 打开文件并获取内容
                with open(file_name_json, 'r') as fp:
                    sub_json = json.load(fp)  # 加载json文件
                    
                    if _del:  # 删除首尾
                        tmp_list = sub_json['lanes']
                        for i in range(len(tmp_list)):
                            tmp_list[i][0] = -1
                            tmp_list[i][-1] = -1
                            tmp_list[i] = [int(x) for x in tmp_list[i]]
                        sub_json['lanes'] = tmp_list
                    #####################
                    json_list.append(sub_json)

        all_json = {}
        all_json['datas'] = json_list

        # 保存最后的json
        for i in range(len(json_list)):
            if i == 0:
                json.dump(json_list[i], open(all_json_path, 'w'), cls=MyEncoder)
                with open(all_json_path, 'a') as f:
                    f.write('\n')
            else:  # 追加json
                json.dump(json_list[i], open(all_json_path, 'a'), cls=MyEncoder)
                with open(all_json_path, 'a') as f:
                    f.write('\n')
        # json.dump(json_list, open(all_json_path, 'w'), indent=4, cls=MyEncoder)
        print('保存成功！文件名：{}'.format(v_pre_name+'.json'))

    def get_curve(self, point_datas, step=1):
        """ 样条插值 """
        x = []
        y = []
        curve_y = np.arange(0, self.now_image.shape[0], step)
        for point in point_datas:
            if point:
                y.append(point[0])
                x.append(point[1])
        if len(x) <= 2:
            return None, None
        x = np.array(x)
        y = np.array(y)
        ipo = spi.splrep(x, y, k=2)
        curve_x = spi.splev(curve_y, ipo)
        return curve_x, curve_y

    def Draw_EditMode(self, im):
        """ 绘制 """
        if self.edit_mode in ['Point', 'Line']:
            if self.line_idx >= 0 and self.line_idx <= 3:
                if self.point_idx >= 0 and self.point_idx <= 4 and self.point_list[self.point_idx]:
                    self.data['Line'][self.line_idx][self.point_idx] = self.point_list[self.point_idx]

        if self.edit_mode in ['Point', 'Line', 'Eraser']:
            if self.edit_mode == 'Eraser':
                x, y = self.now_position
                if self.drawing:
                    for i in range((self.erase_size-1)*2+1):
                        for j in range((self.erase_size-1)*2+1):
                            nx = x+i-self.erase_size+1
                            ny = y+j-self.erase_size+1
                            if nx < 0 or nx >= self.m_erase.shape[0]:
                                continue
                            if ny < 0 or ny >= self.m_erase.shape[1]:
                                continue
                            self.m_erase[nx][ny] = 1
                            self.draw_point(im, (nx, ny), self.color[4], 2)
            elif self.edit_mode in ['Point', 'Line']:
                for idx in range(4):
                    # if ytcz: 样条插值
                    if self.curve_mode:                    
                        curve_x, curve_y = self.get_curve(self.data['Line'][idx])
                        if curve_x is not None:
                            for i in range(len(curve_x)):
                                self.draw_point(im, (int(curve_x[i]), int(curve_y[i])), self.color[idx], size=1)
                    p_list = self.data['Line'][idx]
                    for point in p_list:
                        if point is not None:
                            self.draw_point(im, point, self.color[idx])

    def show_normal_message(self):
        os.system('cls' if os.name == 'nt' else 'clear')
        print("""功能：
        p: play, 播放模式
        a: play, 播放模式（方便单手操作）
        e: edit, 编辑模式
        m: merge, 合并json文件
        q: quit, 退出程序
        ----------------------------------------------------------
        """)

    def run_normal(self):        
        if self.last_im_success is None:
            success, im_origin = self.cap.read()
            if success:
                self.last_im_success = im_origin.copy()
            elif self.last_im_success is None:
                print('无法获取图片')
                self.frame_mode = 'Exit'
                return
        assert self.last_im_success is not None
        self.now_image = self.last_im_success.copy()
        cv2.imshow("door-detection", self.now_image)

        key = cv2.waitKey(1)
        if key == ord('p') or key == ord('a'):     # 播放模式
            self.show_auto_play_message()
            self.frame_mode = 'AutoPlay'
        elif key == ord('e'):  # 编辑模式
            self.show_edit_message()
            self.frame_mode = 'Edit'
            self.edit_mode = 'Line'
            self.line_idx = -1
        elif key == ord('m'):  # 编辑模式
            self.merge_json()
        elif key == ord('q'):
            self.frame_mode = 'Exit'

    def show_auto_play_message(self):
        os.system('cls' if os.name == 'nt' else 'clear')
        print("""当前是播放模式。功能：
        p: play/pause, 切换自动播放
        f: forward, 向前一帧，并停止自动播放
        b: backward, 回退一帧，并停止自动播放
        F: forward 10, 向前10帧，并停止自动播放
        g: forward 100, 向前100帧，并停止自动播放
        G: forward 1000, 向前1000帧，并停止自动播放
        B: backward 10, 回退10帧，并停止自动播放
        h: home, 返回第一帧
        q: quit, 退出播放模式
        ----------------------------------------------------------
        """)

    def read_image(self):
        time.sleep(0.05)
        success, im_origin = self.cap.read()
        if success:
            self.last_im_success = im_origin.copy()
        elif self.last_im_success is None:
            print('无法加载图像')

    def run_auto_play(self, video_paused):
        if not video_paused:
            self.read_image()
        assert self.last_im_success is not None
        self.now_image = self.last_im_success.copy()
        cv2.imshow("door-detection", self.now_image)
        
        key = cv2.waitKey(1)
        if key == ord('q'):     # 退出播放模式
            self.show_normal_message()
            self.frame_mode = 'Normal'
        elif key == ord('f'):   # 单帧步进
            video_paused = True
            self.read_image()
        elif key == ord('F'):  # 单帧步进+10
            video_paused = True
            pos_frame = self.cap.get(cv2.CAP_PROP_POS_FRAMES)
            self.cap.set(cv2.CAP_PROP_POS_FRAMES, pos_frame+10)
            self.read_image()
        elif key == ord('g'):  # 单帧步进+100
            video_paused = True
            pos_frame = self.cap.get(cv2.CAP_PROP_POS_FRAMES)
            self.cap.set(cv2.CAP_PROP_POS_FRAMES, pos_frame + 100)
            self.read_image()
        elif key == ord('G'):  # 单帧步进+1000
            video_paused = True
            pos_frame = self.cap.get(cv2.CAP_PROP_POS_FRAMES)
            self.cap.set(cv2.CAP_PROP_POS_FRAMES, pos_frame + 1000)
            self.read_image()
        elif key == ord('b'):  # 单帧步进-2
            video_paused = True
            pos_frame = self.cap.get(cv2.CAP_PROP_POS_FRAMES)
            if pos_frame>=2:
                self.cap.set(cv2.CAP_PROP_POS_FRAMES, pos_frame-2)
            self.read_image()
        elif key == ord('B'):  # 单帧步进-10
            video_paused = True
            pos_frame = self.cap.get(cv2.CAP_PROP_POS_FRAMES)
            if pos_frame >= 10:
                self.cap.set(cv2.CAP_PROP_POS_FRAMES, pos_frame - 10)
            self.read_image()
        elif key == ord('p'):   # 视频模式和单帧模式切换
            video_paused = False if video_paused else True
        elif key == ord('h'):  # 回到开始 =-1
            self.cap.set(cv2.CAP_PROP_POS_FRAMES, -1)
            video_paused = True
            self.read_image()
        return video_paused

    def show_edit_message(self):
        os.system('cls' if os.name == 'nt' else 'clear')
        pos_frame = self.cap.get(cv2.CAP_PROP_POS_FRAMES)
        print("""当前是编辑模式，第{}帧。功能：
        1: 绘制第1条线
        2: 绘制第2条线
        3: 绘制第3条线
        4: 绘制第4条线
        !: 清空当前绘制的线
        c: 生成线（再按一次返回）
        m: 橡皮檫模式
        q: quit, 退出编辑模式
        ----------------------------------------------------------
        """.format(pos_frame))

    def run_edit(self):
        if self.edit_mode == 'Eraser':
            self.run_edit_eraser()
        elif self.edit_mode == 'Point':
            self.run_edit_point()
        elif self.edit_mode == 'Line':
            assert self.last_im_success is not None
            self.now_image = self.last_im_success.copy()
            self.Draw_EditMode(self.now_image)

            cv2.imshow("door-detection", self.now_image)

            key = cv2.waitKey(1)
            if key == ord('q'):     # 退出编辑模式
                self.show_normal_message()
                self.frame_mode = 'Normal'
            elif key == ord('1'):  # 画门竖线1
                self.edit_mode = 'Point'
                self.line_idx = 0
                self.point_idx = 0
                self.curve_mode = False
                self.bind_mode = False
                self.show_edit_point_message()
            elif key == ord('2'):  # 画门竖线2
                self.edit_mode = 'Point'
                self.line_idx = 1
                self.point_idx = 0
                self.curve_mode = False
                self.bind_mode = False
                self.show_edit_point_message()
            elif key == ord('3'):  # 画门竖线3
                self.edit_mode = 'Point'
                self.line_idx = 2
                self.point_idx = 0
                self.curve_mode = False
                self.bind_mode = False
                self.show_edit_point_message()
            elif key == ord('4'):  # 画门竖线4
                self.edit_mode = 'Point'
                self.line_idx = 3
                self.point_idx = 0
                self.curve_mode = False
                self.bind_mode = False
                self.show_edit_point_message()
            elif key == ord('c'):
                self.curve_mode = False if self.curve_mode else True
            elif key == ord('!'):
                self.data['Line'] = [[None, None, None, None, None],
                        [None, None, None, None, None],
                        [None, None, None, None, None],
                        [None, None, None, None, None]]
            elif key == ord('m'):
                self.edit_mode = 'Eraser'
                self.curve_mode = True
                h, w = self.now_image.shape[:2]
                self.m_erase = np.zeros((w, h), np.uint8)
                self.show_edit_eraser_message()
        
    def show_edit_eraser_message(self):
        os.system('cls' if os.name == 'nt' else 'clear')
        print("""当前是编辑模式，笔刷大小：{}（默认 10）。笔刷功能：
        v: 保存数据
        +: 笔刷大小+5
        -: 笔刷大小-5
        r: 笔刷大小重置
        q: quit, 退出笔刷功能

        说明：使用鼠标在图上擦除部分区域
        ----------------------------------------------------------
        """.format(self.erase_size))

    def run_edit_eraser(self):
        self.Draw_EditMode(self.now_image)
        cv2.imshow("door-detection", self.now_image)
        key = cv2.waitKey(1)
        if key == ord('q'):     # 退出笔刷
            self.show_edit_message()
            self.edit_mode = 'Line'
        elif key == ord('+'):
            self.erase_size += 5
            self.show_edit_eraser_message()
        elif key == ord('-'):
            self.erase_size -= 5
            self.show_edit_eraser_message()
        elif key == ord('r'):
            self.erase_size = 10
            self.show_edit_eraser_message()
        elif key == ord('v'):   # 保存数据
            # save the data
            pos_frame = self.cap.get(cv2.CAP_PROP_POS_FRAMES)
            self.save_data(pos_frame)

    def show_edit_point_message(self):
        os.system('cls' if os.name == 'nt' else 'clear')
        print("""当前是编辑模式。正在绘制第{}条线。画点功能：
        1: 绘制第1个点
        2: 绘制第2个点
        3: 绘制第3个点
        4: 绘制第4个点
        5: 绘制第5个点
        v: 开启&关闭点绑定(当前{}开启)
        q: quit, 保存并退出画点功能

        说明：1. 可以画1-5个点，但是至少3个点才能使用二次样条插值生成线
        ----------------------------------------------------------
        """.format(self.line_idx+1, "已" if self.bind_mode else "未"))

    def run_edit_point(self):
        assert self.last_im_success is not None
        self.now_image = self.last_im_success.copy()

        self.Draw_EditMode(self.now_image)

        cv2.imshow("door-detection", self.now_image)

        key = cv2.waitKey(1)
        if key == ord('1'):
            self.point_idx = 0
            self.curve_mode = False
        elif key == ord('2'):
            self.point_idx = 1
            self.curve_mode = False
        elif key == ord('3'):
            self.point_idx = 2
            self.curve_mode = False
        elif key == ord('4'):
            self.point_idx = 3
            self.curve_mode = False
        elif key == ord('5'):
            self.point_idx = 4
            self.curve_mode = False
        elif key == ord('v'):
            self.bind_mode = False if self.bind_mode else True
            self.show_edit_point_message()
        elif key == ord('q'):     # 退出画点
            self.edit_mode = 'Line'
            self.bind_mode = False
            self.point_idx = -1
            self.show_edit_message()

    def run(self):
        self.now_image = np.zeros((720, 1280, 3), np.uint8)
        cv2.namedWindow("door-detection", 0)
        cv2.setMouseCallback("door-detection", self.label_mask, self.now_image)
        self.cap = cv2.VideoCapture(self.video_path)
        self.last_im_success = None
        self.frame_mode = 'Normal'
        self.show_normal_message()
        video_paused = True
        while self.cap.isOpened():
            if self.frame_mode == 'Normal':
                self.run_normal()
            elif self.frame_mode == 'AutoPlay':
                video_paused = self.run_auto_play(video_paused)
            elif self.frame_mode == 'Edit':
                self.run_edit()
            elif self.frame_mode == 'Exit':
                break

        self.cap.release()
        cv2.destroyAllWindows()
        pass

if __name__=='__main__':
    # label_door = LabelDoor(r'D:\PythonWorkroom\datasets\Elevator\School_With_Tags/20210222_044216_rot180.mp4')   # 未标注
    # label_door = LabelDoor(r'video/20210222_050945.mp4')  # 已标注
    # label_door = LabelDoor(r'D:\PythonWorkroom\datasets\Elevator\School_With_Tags/20210222_050945.mp4')  # 已标注
    label_door = LabelDoor(r'D:\PythonWorkroom\datasets\Elevator\School_With_Tags/20210222_052349.mp4')   # 已标注
    # label_door = LabelDoor(r'D:\PythonWorkroom\datasets\Elevator\School_With_Tags/20210222_053738.mp4')   # 已标注
    # label_door = LabelDoor(r'D:\PythonWorkroom\datasets\Elevator\School_With_Tags/20210222_054658.mp4')   # 不标注了，与053738是同一部电梯
    # label_door = LabelDoor(r'D:\PythonWorkroom\datasets\Elevator\School_With_Tags/20210222_230723_rot180.mp4')   # 已标注
    # label_door = LabelDoor(r'D:\PythonWorkroom\datasets\Elevator\School_With_Tags/20210222_231810_rot180.mp4')   # 已标注
    label_door.run()